# -*- coding: utf-8 -*-
# Created by 'Zhou Bingbing'  on 2019/7/30


import numpy as np


class HMM:
    def forward(self, Q, V, A, B, O, PI):
        """
        前向算法
        Q：状态序列
        V：观测集合
        A：状态转移概率矩阵
        B：生成概率矩阵
        O：观测序列
        PI:状态概率向量
        """
        N = len(Q) #状态
        M = len(O) #观测序列
        # alphas是前向概率矩阵，每列是某个时刻每个状态的前向概率
        alphas = np.zeros((N, M))
        T = M  # 时刻 观察序列
        for t in range(T):
            # 当前时刻观测值对应的索引
            index_O = V.index(O[t])
            for i in range(N): #状态序列
                if t == 0:
                    alphas[i][t] = PI[0][i] * B[i][index_O]
                else:
                    # 这里使用点积和使用矩阵相乘后再sum没区别
                    alphas[i][t] = np.dot([alpha[t - 1] for alpha in alphas], [a[i] for a in A]) * B[i][index_O]
        # 将前向概率的最后一列的概率相加就是观测序列生成的概率
        P = np.sum([alpha[M - 1] for alpha in alphas])
        print('alphas :')
        print(alphas)
        print('P=%5f' % P)

    def backward(self, Q, V, A, B, O, PI):
        """后向算法"""
        N = len(Q)
        M = len(O)
        # betas 是后向概率矩阵，第M列的概率都为1
        betas = np.zeros((N, M))
        for i in range(N):
            betas[i][M - 1] = 1
        # 倒着，从后往前递推
        for t in range(M - 2, -1, -1):
            index_O = V.index(O[t])
            for i in range(N):
                betas[i][t] = np.dot(np.multiply(A[i], [b[index_O] for b in B]), [beta[t + 1] for beta in betas])
                realT = t + 1
                realI = i + 1
        index_O = V.index(O[0])
        P = np.dot(np.multiply(PI, [b[index_O] for b in B]), [beta[0] for beta in betas])
        print("P(O|lambda)=", end="")
        for i in range(N):
            print("%.1f*%.1f*%.5f+" % (PI[0][i], B[i][index_O], betas[i][0]), end="")
        print("0=%f" % P)
        print(betas)

    def viterbi(self, Q, V, A, B, O, PI):
        N = len(Q) #状态序列
        M = len(O)-1 #观察序列
        # deltas是记录每个时刻，每个状态的最优路径的概率
        deltas = np.zeros((N, M))
        # psis是记录每个时刻每个状态的前一个时刻的概率最大的路径的节点（psia的记录是比deltas晚一个时刻的）
        psis = np.zeros((N, M))
        # I是记录最优路径的
        I = np.zeros((1,M))
        for t in range(M): #观察序列

            index_O = V.index(O[t]) #获取我们的观察序列的所对应的索引
            for i in range(N):#状态

                if t == 0:
                    deltas[i][t] = PI[0][i] * B[i][index_O]
                    psis[i][t] = 0
                else:
                    deltas[i][t] = np.max(np.multiply([delta[t - 1] for delta in deltas], [a[i] for a in A])) * B[i][
                        index_O]
                    psis[i][t] = np.argmax(np.multiply([delta[t - 1] for delta in deltas], [a[i] for a in A]))
        # 直接取最后一列最大概率对应的状态为最优路径最后的节点
        I[0][M - 1] = np.argmax([delta[M - 1] for delta in deltas])
        print([delta[M - 1] for delta in deltas])
        print(I)
        for t in range(M - 2, -1, -1):
            # psis要晚一个时间步，起始将最后那个状态对应在psis那行直接取出就是最后的结果
            # 但是那样体现不出回溯，下面这种每次取上一个最优路径点对应的上一个最优路径点
            I[0][t] = psis[int(I[0][t + 1])][t + 1]
        print(I)


if __name__ == '__main__':
    """
            前向算法
            Q：状态序列
            V：观测集合
            A：状态转移概率矩阵
            B：生成概率矩阵
            O：观测序列
            PI:状态概率向量
            """
    Q = [1, 2, 3] #状态序列
    V = ['红', '白'] #观测集合
    A = [[0.5, 0.2, 0.3], [0.3, 0.5, 0.2], [0.2, 0.3, 0.5]] #状态转换概率
    B = [[0.5, 0.5], [0.4, 0.6], [0.7, 0.3]] #状态发射概率
    # O = ['红', '白', '红', '红', '白', '红', '白', '白']
    O = ['红', '白', '红', '白']  #观察序列
    PI = [[0.2, 0.4, 0.4]] #初始状态转移概率
    Hmm = HMM()
    Hmm.forward(Q, V, A, B, O, PI)
    # Hmm.backward(Q, V, A, B, O, PI)
    Hmm.viterbi(Q, V, A, B, O, PI)
